import torch

# root用户从本地加载
model_path = '/root/.cache/torch/hub/huggingface_pytorch-transformers_main'

# 直接使用预训练的bert中文模型
model_name = 'bert-base-chinese'
# 通过torch.hub获得已经训练好的bert-base-chinese模型
model =  torch.hub.load(model_path, 'model', model_name, source='local')
# 获得对应的字符映射器，它将把中文的每个字映射成一个数字
tokenizer = torch.hub.load(model_path, 'tokenizer', model_name, source='local')

def get_bert_encode_for_single(text:str):
    """
    把传入的中文文本，使用bert-base-chinese模型转成字向量
    :param text: 输入的文本
    :return: 转成的字向量
    """

    indexed_token = tokenizer.encode(text)[1:-1]
    token_tensor = torch.LongTensor([indexed_token])

    with torch.no_grad():
        encode_layer = model(token_tensor)

    return encode_layer[0]


if __name__ == '__main__':

    text = "你好，周杰伦"

    outputs = get_bert_encode_for_single(text)
    print(outputs)
    # print(outputs.shape)
